// parser/parser.cpp
// Written by, Sohail Qayum Malik[sqm@hackers.pk]

/* For debuging purposes, should be gone when done */
#include <iostream>
/*  For debugging purposes ends here */

#include <WINDOWS.H>

//#include "../args.h"
#include "parser.h"

using namespace::std;

int parser(const char *filename, int (*callback)(int, ...), unsigned int cmd) throw(args_exception) {

   int ret = 0;    
   void *ptr = NULL;
   fstream PEfile; 	    
   IMAGE_DOS_HEADER  dosheader;

   //IMAGE_NT_HEADERS, which is defined in WINNT.H. This structure is composed of a DWORD and two substructures and is laid out as follows
   /* DWORD Signature;
      IMAGE_FILE_HEADER FileHeader;
      IMAGE_OPTIONAL_HEADER32/64 OptionalHeader; */
#ifdef _WIN64
   IMAGE_NT_HEADERS64 ntheader;
#else   
   IMAGE_NT_HEADERS32 ntheader;
#endif   
         
   PEfile.open(filename, fstream::in);
   
   if ( PEfile.fail() ) 
	  throw( args_exception("Unable to open file " + string(filename)) );
	    
   PEfile.read( (char *)&dosheader, sizeof(IMAGE_DOS_HEADER) ); 
   if ( PEfile.fail() || PEfile.eof() ) {
  
	  PEfile.close();
	  throw( args_exception("Failed to read IMAGE_DOS_HEADER from file " + string(filename)) );
   }
   
   PEfile.seekg( dosheader.e_lfanew, fstream::beg );
   if ( PEfile.fail() ) {
   
      PEfile.close();
      throw( args_exception("seekg() operation failed with fstream::beg flag on file " + string(filename)) ); 	  
   }
   
   PEfile.read( (char *)&ntheader, sizeof(ntheader) );
   if ( PEfile.fail() || PEfile.eof() ) {
   
      PEfile.close();
	  
#ifdef _WIN64   
      throw( args_exception("Failed t read IMAGE_NT_HEADERS64 from file " +
	  string(filename)) );
#else
	  throw( args_exception("Failed to read IMAGE_NT_HEADERS32 from file " +
	  string(filename)) );  
#endif	  
      PEfile.close(); 
   }  
          
   if ( dosheader.e_magic != IMAGE_DOS_SIGNATURE )
      throw( args_exception("file " + string(filename) +" has no IMAGE_DOS_SIGNATURE") );  
/*   else 
      std::cout<<std::hex<<dosheader.e_magic<<std::endl;   */
 	  
   if ( ntheader.Signature != IMAGE_NT_SIGNATURE )
      throw( args_exception("file " + string(filename) +" has no IMAGE_NT_SIGNATURE") );  	  
/*   else 
      std::cout<<std::hex<<ntheader.Signature<<std::endl;  */ 
	
   switch(cmd) {
   
      case PARSER_DOS_HEADER: 
	  
	     // IMAGE_DOS_HEADER and IMAGE_NT_HEADERS were read and available for use. We'll now place the file pointer just passed the IMAGE_DOS_HEADER + The bytes of DOS stub routine(this routine prints a message when a Windows exe file gets executed in the DOS environment 	
         PEfile.seekg( /* sizeof(IMAGE_DOS_HEADER) */ dosheader.e_cparhdr << 4 , fstream::beg );
		 
		 if ( PEfile.fail() ) {
   
            PEfile.close();
            throw( args_exception("seekg() operation failed with fstream::beg flag on file " + string(filename)) ); 	  
         }	  

         ptr = malloc( dosheader.e_lfanew - sizeof(IMAGE_DOS_HEADER) );
   
         if ( ptr == NULL ) {
   
            PEfile.close();
            throw( args_exception("Failed to allocate memory") );
         }
		    
         PEfile.read( (char *)ptr, dosheader.e_lfanew - sizeof(IMAGE_DOS_HEADER)); 
		 
         if ( PEfile.fail() || PEfile.eof() ) {
  
            free( ptr );
	        PEfile.close();
	        throw( args_exception("Failed to read DOS stub from file " + string(filename)) );
         }
	  
	     // A place to put the uggliest code in the world
         ret = callback(4, &dosheader, &ntheader, (size_t)(dosheader.e_lfanew - sizeof(IMAGE_DOS_HEADER) ), ptr); 
      break;
	  case PARSER_IMAGE_NT_HEADERS:

         // IMAGE_DOS_HEADER and IMAGE_NT_HEADERS were read and available for use. We'll now place the file pointer just passed the IMAGE_DOS_HEADER + The bytes of DOS stub routine(this routine prints a message when a Windows exe file gets executed in the DOS environment 	
   
#ifdef _WIN64
         PEfile.seekg( dosheader.e_lfanew + sizeof(IMAGE_NT_HEADERS64), fstream::beg );  
#else
         PEfile.seekg( dosheader.e_lfanew + sizeof(IMAGE_NT_HEADERS32), fstream::beg );
#endif

         if ( PEfile.fail() ) {
   
            PEfile.close();
            throw( args_exception("seekg() operation failed with fstream::beg flag on file " + string(filename)) ); 	  
         }	  
	  
	     // A place to put uggliest code in the world
	     ret = callback(4, &dosheader, &ntheader, (size_t)0, NULL);
	  break;	  
	  case PARSER_IMAGE_SECTION_HEADERS:
	  
	     //std::cout<<"Number of sections = "<<ntheader.FileHeader.NumberOfSections<<std::endl; 
		 
		// IMAGE_DOS_HEADER and IMAGE_NT_HEADERS were read and available for use. We'll now place the file pointer just passed the IMAGE_DOS_HEADER + The bytes of DOS stub routine(this routine prints a message when a Windows exe file gets executed in the DOS environment 	   
#ifdef _WIN64
         PEfile.seekg( dosheader.e_lfanew + sizeof(IMAGE_NT_HEADERS64), fstream::beg );  
#else
         PEfile.seekg( dosheader.e_lfanew + sizeof(IMAGE_NT_HEADERS32), fstream::beg );
#endif
	  
	     if ( PEfile.fail() ) {
   
            PEfile.close();
            throw( args_exception("seekg() operation failed with fstream::beg flag on file " + string(filename)) ); 	  
         }	  
	  
	     ptr =  malloc(IMAGE_SIZEOF_SECTION_HEADER*ntheader.FileHeader.NumberOfSections);
		 
         if ( ptr == NULL ) {
   
            PEfile.close();
            throw( args_exception("Failed to allocate memory") );
         }


         PEfile.read( (char *)ptr, IMAGE_SIZEOF_SECTION_HEADER*ntheader.FileHeader.NumberOfSections); 
		 
         if ( PEfile.fail() || PEfile.eof() ) {
  
            free( ptr );
	        PEfile.close();
	        throw( args_exception("Failed to read IMAGE_SECTION_HEADERS from file " + string(filename)) );
         }		 
		 	  
	     ret = callback(4, &dosheader, &ntheader, (size_t)ntheader.FileHeader.NumberOfSections, ptr);
	  break;
   }
   
   free(ptr);   
   PEfile.close();
                    	      	  	           	  	  	         
   return ret;
}